#include "med.hh"
#ifdef MEDGAUSS
#include "med_gauss.hh"
#endif
#include "dpm.hh"

#include <boost/lexical_cast.hpp>

template<typename T>
T get_param(const char* arg, std::string emsg)
{
  T val;
  try
  {
    val = boost::lexical_cast<T>(arg);
  } catch (boost::bad_lexical_cast&)
  {
    err_msg(emsg);
    exit(1);
  }
  return val;
}

////////////////////////////////////////////////////////////////
struct prog_args_t
{
  explicit prog_args_t() :
    data_file(""), output("output"), ncols(0), // required args
    tol(1e-3), max_iter(100), sample(100), K(2),// options
    repeat(1), bmin(-2.), bmax(2.), minibatch(10) // 
    {
    }

  std::string data_file;
  std::string output;
  std::vector<size_t> ncols;
  double tol;
  int max_iter;
  int sample;
  int K;
  int repeat;
  double bmin;
  double bmax;
  double minibatch;
};

void print_help(const char* prog, const prog_args_t& args)
{
  std::cerr << prog << " -d data.txt -k max_num_clusters -o out -c c1 -c c2 ..." << std::endl;
  std::cerr << "         -bmin bmin -bmax bmax" << std::endl;
  std::cerr << "         -i iter -t tol -s samples -r repeats" << std::endl;
  std::cerr << "         -minibatch size" << std::endl;
  std::cerr << std::endl;
  std::cerr << " data.txt      expression matrix (n x p tab-separated)" << std::endl;
  std::cerr << " c[N]          # columns of N-th condition" << std::endl;
  std::cerr << "               such that sum_i c[i] = p" << std::endl;
  std::cerr << " max_n_clust   maximum number of clusters" << std::endl;
  std::cerr << " out           output files' header (default: " << args.output << ")" << std::endl;
  std::cerr << " bmin           min value of beta" << std::endl;
  std::cerr << " bmax           max value of beta" << std::endl;  
  std::cerr << " minibatch      size of mini-batch (default: " << args.minibatch
	    << ")" << std::endl;
  std::cerr << std::endl;
}

bool parse_args(const int argc, const char* argv[], prog_args_t& args_out)
{
  for (int j = 1; j < argc; ++j)
  {
    std::string curr = argv[j];

    if (curr == "-d" && ++j < argc)
    {
      args_out.data_file = argv[j];
    }
    else if (curr == "-o" && ++j < argc)
    {
      args_out.output = argv[j];
    }
    else if (curr == "-c" && ++j < argc)
    {
      args_out.ncols.push_back( get_param<size_t>(argv[j], "parsing ncol") );
    }
    else if (curr == "-k" && ++j < argc)
    {
      args_out.K = get_param<int>(argv[j], "parsing K");
    }
    else if (curr == "-i" && ++j < argc)
    {
      args_out.max_iter = get_param<int>(argv[j], "parsing iter");
    }
    else if (curr == "-r" && ++j < argc)
    {
      args_out.repeat = get_param<size_t>(argv[j], "parsing repeat");
    }
    else if (curr == "-t" && ++j < argc)
    {
      args_out.tol = get_param<double>(argv[j], "parsing tol");
    }
    else if (curr == "-bmin" && ++j < argc)
    {
      args_out.bmin = get_param<double>(argv[j], "parsing beta min");
    }
    else if (curr == "-bmax" && ++j < argc)
    {
      args_out.bmax = get_param<double>(argv[j], "parsing beta max");
    }
    else if (curr == "-minibatch" && ++j < argc)
    {
      args_out.minibatch = get_param<size_t>(argv[j], "parsing minibatch");
    }

  }

  if (args_out.tol <= 0)
    return false;
  if (args_out.data_file.size() == 0 || args_out.ncols.size() == 0 )      
    return false;
  if (args_out.bmin >= args_out.bmax)
    return false;
  if (args_out.K < 1)
    return false;
  return true;
}

////////////////////////////////////////////////////////////////

#ifdef MEDGAUSS
typedef func_expr_pairwise_div_gauss_t func_type;
#else
typedef func_expr_pairwise_div_t func_type;
#endif

int
main(const int argc, const char* argv[])
{

  prog_args_t args;

  if( ! parse_args( argc, argv, args ) )
  {
    print_help( argv[0], args );
    exit(1);
  }

  const char* data_file = args.data_file.c_str();
  std::vector<size_t>& ncols = args.ncols;
  const size_t K = args.K;

  expr_tuple_vec_t data_vec;
  if( ! read_expr_tuple( data_file, ncols, data_vec ) )
  {
    TLOG("cannot read data");
    exit(1);
  }

  double max_llik = 0.;
  boost::ptr_vector<func_type> components;

  for (size_t r = 1; r <= args.repeat; ++r)
  {

    components.clear();
    for (size_t k = 0; k < K; ++k)
      components.push_back(new func_type(ncols.size(), args.bmin, args.bmax));
    
    double curr_llik;

#ifdef ONLINE
    curr_llik = fit_dpm_online<expr_tuple_t, func_type>( data_vec, components, args.max_iter, args.sample, args.tol, args.minibatch );      
#else
    curr_llik = fit_dpm<expr_tuple_t, func_type>(data_vec, components, args.max_iter, args.sample, args.tol);
#endif

    if (r == 1 || curr_llik > max_llik)
    {    
      max_llik = curr_llik;
    
      TLOG(" repeat = " << r << " max log-likelihood = " << max_llik);

      // output argmax membership
      std::ofstream ofs((args.output + ".argmax").c_str(), std::ios::out);
      for (size_t i = 0; i < data_vec.size(); ++i)
      {
	const expr_tuple_t& expr = *data_vec.at(i).get();
	size_t argmax = 0;
	double maxval = components[argmax].score(expr);
	for (size_t k = 1; k < K; ++k)
	{
	  double curr = components[k].score(expr);	  
	  if (curr > maxval)
	  {
	    maxval = curr;
	    argmax = k;
	  }
	}
	ofs << i << "\t" << argmax << "\t" << maxval << std::endl;
      }
      ofs.close();

      // output parameters
      for (size_t k = 0; k < K; ++k)
      {
	const func_type& F = components[k];

	std::ofstream ofs_par((args.output + ".beta" +
			       boost::lexical_cast<std::string>(k)).c_str(),
			      std::ios::out);
	
	F.write_beta( ofs_par );
	ofs_par.close();
      }
#ifdef MEDGAUSS
      for (size_t k = 0; k < K; ++k)
      {
	const func_type& F = components[k];
	std::ofstream ofs_par((args.output + ".mu" +
			       boost::lexical_cast<std::string>(k)).c_str(),
			      std::ios::out);
	
	F.write_mu( ofs_par );
	ofs_par.close();
      }      
#endif
    }
  }

  return 0;
}
